10.2 배포와 서빙 - vLLM의 Mamba 지원 현황 및 ONNX/TensorRT 변환 이슈

10.2 배포와 서빙 - vLLM의 Mamba 지원 현황 및 ONNX/TensorRT 변환 이슈

1. 서론: 포스트 트랜스포머 시대의 서빙 패러다임 변화

인공지능, 특히 거대언어모델(Large Language Model, LLM)의 발전사는 지난 수년간 트랜스포머(Transformer) 아키텍처의 독주로 정의되어 왔다. 구글(Google)의 “Attention Is All You Need” 논문 이후, 어텐션 메커니즘은 자연어 처리의 표준이 되었고, 이에 따라 모델을 배포하고 서빙(Serving)하는 모든 인프라스트럭처 역시 트랜스포머의 연산 특성에 맞춰 최적화되었다. NVIDIA의 GPU 아키텍처부터 CUDA 커널, 그리고 vLLM이나 TensorRT-LLM과 같은 고수준 서빙 엔진에 이르기까지, 모든 기술 스택은 O(N^2)의 복잡도를 가지는 행렬 곱 연산과 KV 캐시(Key-Value Cache) 관리에 초점을 맞추어 진화해 온 것이다.

그러나 2023년 말부터 본격화된 상태 공간 모델(State Space Model, SSM), 그중에서도 Mamba 아키텍처의 부상은 이러한 서빙 생태계에 근본적인 질문을 던지고 있다. Mamba는 입력 시퀀스 길이에 대해 선형적인(O(N)) 연산 복잡도와 상수(O(1)) 메모리 요구량을 통해 트랜스포머의 고질적인 병목인 긴 문맥(Long Context) 처리 비용 문제를 해결했다.1 또한, 어텐션과 SSM을 결합하여 추론 능력과 효율성의 균형을 꾀하는 Jamba와 같은 하이브리드(Hybrid) 아키텍처의 등장은 모델의 다양성을 폭발적으로 증가시켰다.1

이러한 모델 아키텍처의 다변화는 서빙 엔지니어링 관점에서 “이중적 과제“를 안겨준다. 기존 트랜스포머를 위해 구축된 고도로 최적화된 파이프라인(예: PagedAttention)을 유지하면서도, 전혀 다른 메모리 접근 패턴과 연산 특성을 가진 SSM을 통합해야 하기 때문이다. 본 보고서에서는 현재 가장 널리 사용되는 오픈소스 서빙 엔진인 vLLM의 Mamba 및 하이브리드 모델 지원 현황을 아키텍처 레벨에서 심층 분석하고, 프로덕션 배포의 핵심 관문인 ONNX 및 TensorRT로의 변환 과정에서 발생하는 기술적 난제와 그 해결 방안을 상세히 기술한다. 특히 vLLM v1 아키텍처로의 전환 과정에서 드러난 메모리 관리의 기술적 부채와 해결 전략, 그리고 하드웨어 가속을 위한 TensorRT-LLM의 최적화 기법을 면밀히 비교 분석함으로써, 차세대 모델 서빙을 위한 실질적인 가이드라인을 제시하고자 한다.

2. Mamba 및 하이브리드 아키텍처의 서빙 복잡성 분석

Mamba와 같은 SSM 기반 모델을 효율적으로 서빙하기 위해서는 먼저 이들이 기존 트랜스포머와 어떻게 다른지, 그리고 그 차이가 시스템 리소스 관리에 어떤 영향을 미치는지 이해해야 한다.

2.1 순환 상태(Recurrent State) vs KV 캐시(KV Cache)

트랜스포머 모델 서빙의 핵심은 KV 캐시 관리이다. 디코딩 단계에서 이전에 생성된 모든 토큰의 Key와 Value 벡터를 GPU 메모리에 저장해야 하며, 이는 시퀀스 길이가 길어질수록 선형적으로 증가하여 결국 메모리 병목을 유발한다. vLLM의 PagedAttention은 바로 이 KV 캐시의 메모리 단편화(Fragmentation)를 해결하기 위해 운영체제의 페이징 기법을 차용한 기술이다.4

반면, Mamba는 ’캐시’가 아닌 ’상태(State)’를 관리한다. Mamba 레이어는 이전의 모든 정보를 고정된 크기의 은닉 상태(Hidden State)로 압축하여 전달한다. 이는 이론적으로 무한한 길이의 시퀀스를 처리하더라도 추론 시점의 메모리 사용량이 일정함을 의미한다. 그러나 서빙 엔진 입장에서는 이것이 새로운 도전이 된다.

  • 메모리 할당의 불일치: vLLM과 같은 엔진은 동적으로 커지는 KV 캐시를 블록 단위로 할당하는 데 최적화되어 있다. 고정된 크기지만, 모델의 차원(d_{model})과 상태 확장 계수(Expansion Factor)에 따라 매우 커질 수 있는 Mamba의 상태 텐서를 기존 블록 관리 시스템에 억지로 끼워 맞추는 것은 비효율을 초래한다.5
  • 연산 커널의 이질성: 트랜스포머는 행렬 곱(GEMM) 위주의 연산이지만, Mamba는 ’선택적 스캔(Selective Scan)’이라는 순차적 의존성이 강한 연산을 수행한다. 이는 병렬화가 어렵고, 메모리 대역폭보다는 SRAM(Static RAM)과 HBM(High Bandwidth Memory) 사이의 데이터 이동 효율성이 성능을 좌우한다.

2.2 하이브리드 아키텍처(Jamba)의 딜레마

AI21 Labs의 Jamba 모델은 트랜스포머 레이어와 Mamba 레이어를 교차 배치하고, 일부 레이어에는 MoE(Mixture of Experts)를 적용한 하이브리드 구조이다.3 이러한 구조는 서빙 시스템의 복잡도를 기하급수적으로 높인다.

  • 이질적 메모리 요구량: 하나의 요청(Request)을 처리하는 동안, 트랜스포머 레이어는 KV 캐시를 위해 지속적으로 추가 메모리를 할당받아야 하는 반면, Mamba 레이어는 고정된 상태 메모리만 점유하면 된다. 단일 메모리 풀(Memory Pool)에서 이 두 가지 서로 다른 요구사항을 동시에 만족시키려 할 때, 할당자(Allocator)의 로직은 복잡해질 수밖에 없다.
  • MoE와 통신 오버헤드: Jamba의 MoE 구조는 텐서 병렬화(Tensor Parallelism)나 전문가 병렬화(Expert Parallelism)를 필요로 한다. Mamba 레이어의 상태 관리와 MoE의 라우팅(Routing) 통신이 결합되면, GPU 간 동기화 비용이 전체 추론 레이턴시에 큰 영향을 미치게 된다.

3. vLLM의 Mamba 통합: v0에서 v1으로의 진화

vLLM 프로젝트는 Mamba 및 하이브리드 모델을 지원하기 위해 아키텍처를 대대적으로 수선해 왔다. 이 과정은 초기 실험적 지원(v0)에서 “일등 시민(First-class Citizen)“으로서의 완전한 통합(v1)으로 나아가는 과정으로 해석할 수 있다.

3.1 vLLM v0: MambaCacheManager와 구조적 한계

초기 vLLM(v0.x 버전)에서의 Mamba 지원은 기존 트랜스포머 파이프라인에 대한 일종의 ‘패치(Patch)’ 형태로 구현되었다. 개발진은 Mamba의 상태 관리를 위해 MambaCacheManager라는 별도의 클래스를 도입했으나, 이는 근본적인 한계를 내포하고 있었다.6

3.1.1 구현 방식과 문제점

MambaCacheManagermodel_executor/models/mamba_cache.py에 정의되어 있으며, Mamba 모델의 상태 텐서를 관리하는 역할을 담당했다. 그러나 이 관리자는 vLLM의 핵심인 **블록 관리자(Block Manager)**와 긴밀하게 통합되지 못했다.

  • 프리픽스 캐싱(Prefix Caching) 불가: vLLM의 강력한 기능인 프리픽스 캐싱은 KV 캐시 블록의 해시값을 비교하여 중복 계산을 방지한다. 그러나 MambaCacheManager는 상태 텐서를 블록 단위로 관리하고 추적하는 메커니즘이 부재했거나 기존 로직과 호환되지 않아, Mamba 모델에서는 이 기능을 사용할 수 없었다. 이는 챗봇이나 에이전트와 같이 시스템 프롬프트를 공유하는 애플리케이션에서 심각한 성능 저하를 의미한다.5
  • 메모리 단편화 및 낭비: 하이브리드 모델의 경우, 어텐션 레이어와 Mamba 레이어가 혼재되어 있다. v0 아키텍처에서는 모든 레이어에 대해 동일한 크기의 캐시 공간을 할당하려는 경향이 있었다. 예를 들어, 어텐션 레이어의 토큰당 메모리 요구량이 작더라도, Mamba 레이어의 큰 상태 크기에 맞춰 과도한 메모리를 할당해야 했고, 이는 전체 GPU 메모리의 낭비로 이어졌다. 일부 시나리오에서는 메모리 낭비율이 79.6%에 달한다는 보고도 있었다.6

3.2 vLLM v1: SingleTypeKVCacheManager와 통합 아키텍처

vLLM v1은 이러한 문제를 해결하기 위해 메모리 관리 시스템을 재설계하였다. 핵심은 모든 유형의 어텐션 및 상태 메커니즘을 포괄할 수 있는 **통합 메모리 할당자(Unified Allocator)**의 도입이다.8

3.2.1 계층별 캐시 관리 (Per-layer Cache Management)

v1 아키텍처에서는 더 이상 모델 전체를 단일한 캐시 관리자가 담당하지 않는다. 대신 레이어의 유형(Attention Type)에 따라 그룹을 나누고, 각 그룹을 SingleTypeKVCacheManager가 전담 관리한다.9

  • 구조: KVCacheCoordinator가 최상위에서 조율을 담당하며, 하위에 여러 개의 SingleTypeKVCacheManager 인스턴스를 거느린다. 예를 들어 Jamba 모델이라면, 트랜스포머 레이어 그룹을 위한 매니저와 Mamba 레이어 그룹을 위한 매니저가 각각 존재한다.
  • 이점: 각 매니저는 자신이 담당하는 레이어의 특성(KV 크기, 상태 크기, 슬라이딩 윈도우 여부 등)에 맞춰 최적의 블록 할당 정책을 수행할 수 있다. 이를 통해 이질적인 레이어가 혼재된 하이브리드 모델에서도 메모리 단편화를 최소화할 수 있다.

3.2.2 하이브리드 모델을 위한 메모리 패딩 전략

Jamba와 같은 모델을 지원하기 위한 구체적인 메모리 할당 알고리즘은 다음과 같다. Mamba 레이어의 상태 크기(\text{state\_size})가 어텐션 레이어의 KV 크기(\text{kv\_hidden\_size})보다 클 때, vLLM은 단일 메모리 풀을 공유하기 위해 어텐션 레이어의 블록 크기(\text{block\_size})를 조정한다.9

\text{block\_size} \times \text{kv\_hidden\_size}_{\text{att}} \ge \text{state\_size}_{\text{mamba}}

즉, 어텐션 레이어의 한 블록이 담을 수 있는 데이터의 총량이 Mamba 레이어의 한 상태를 담기에 충분하도록 블록 크기를 키우는 것이다.

  • 부작용: 이 방식은 Mamba 상태 크기가 매우 클 경우, 어텐션 레이어의 블록 크기를 지나치게 크게(예: 400 이상) 설정하게 만든다. 블록 크기가 커지면 미세한 메모리 관리가 어려워져 다시 단편화가 발생할 수 있다.
  • 개선 방향: 현재 vLLM 팀은 단일 블록이 아닌, 여러 어텐션 레이어의 블록을 합쳐서 Mamba 상태를 저장하는 방식(\text{block\_size} \times \text{kv\_hidden\_size}_{\text{att}} \times \text{num\_attn\_layers} \ge \text{state\_size}_{\text{mamba}})을 개발 중이다. 이는 “가상화된” 메모리 공간을 통해 물리적 메모리 제약을 우회하려는 시도이다.9

3. Jamba 및 Mamba-2 지원의 현재 상태

  • Jamba: JambaForCausalLM 클래스를 통해 지원되며, 전문가(Expert) 레이어에 대한 INT8 양자화 및 텐서 병렬화를 지원한다. Jamba 1.5 Large/Mini 모델 역시 vLLM을 통해 서빙 가능하다.10
  • Mamba-2: 구조적 상태 공간 이원성(SSD)을 특징으로 하는 Mamba-2 역시 지원된다. 다만, Mamba-2의 상태 관리 효율성이 개선되었음에도 불구하고, vLLM의 PagedAttention 기반 메모리 풀과의 완전한 통합은 여전히 진행 중인 과제이다. 특히 텐서 병렬화 지원이 최근에야 추가되는 등(TensorRT-LLM 대비), 대규모 배포를 위한 기능은 점진적으로 강화되고 있다.12

4. 기능적 제약 사항 상세

vLLM에서 Mamba 모델을 사용할 때 사용자는 다음과 같은 제약 사항을 인지해야 한다.

  • 프리픽스 캐싱 미지원 (Work in Progress): 앞서 언급했듯, Mamba 모델에 대한 프리픽스 캐싱은 아직 지원되지 않는다. Mamba의 순환 상태는 트랜스포머의 KV 캐시처럼 토큰 단위로 명확하게 분리되지 않고, 이전 문맥 전체가 압축되어 섞여 있는 형태이기 때문에, 특정 지점에서 상태를 잘라내어 재사용(Caching)하는 로직을 구현하기가 수학적으로 훨씬 복잡하다.7
  • 투기적 디코딩의 불안정성: Mamba 모델을 드래프트(Draft) 모델로 사용하거나 타겟 모델로 사용하는 투기적 디코딩 기능은 버그가 보고되거나 실험적인 단계에 머물러 있다.13

10.2.4 ONNX 변환: 표준화의 장벽과 기술적 해결

ONNX(Open Neural Network Exchange)는 특정 딥러닝 프레임워크에 종속되지 않고 모델을 배포하기 위한 산업 표준이다. 그러나 Mamba와 같은 최신 아키텍처, 특히 사용자 정의 CUDA 커널에 깊게 의존하는 모델을 ONNX로 변환하는 것은 “표준“과 “최적화” 사이의 충돌을 야기한다.

1. Selective Scan 연산자의 블랙박스 문제

Mamba 아키텍처의 심장은 selective_scan 연산이다. PyTorch 구현체에서는 성능을 위해 이 부분을 Python이 아닌 Triton이나 CUDA C++로 작성된 커널을 호출하여 처리한다. torch.onnx.export 함수는 PyTorch의 연산 그래프를 추적(Tracing)하여 ONNX 그래프를 생성하는데, 이 과정에서 컴파일된 커널은 내부가 보이지 않는 ’블랙박스’로 취급된다.14

  • 증상: 변환 시도 시 “Unknown operator” 에러가 발생하거나, 그래프가 끊어진 상태로 생성된다. 또는 TracerWarning과 함께 동적 제어 흐름(Dynamic Control Flow)이 상수로 박제되어, 특정 입력 길이에만 동작하는 기형적인 모델이 만들어지기도 한다.15
  • 원인: ONNX 표준 연산자 세트(Opset)에는 SelectiveScan이라는 연산자가 존재하지 않는다. 따라서 익스포터는 이를 어떻게 표현해야 할지 알 수 없으며, 내부 로직을 기본 연산자(Add, Mul 등)로 풀어서 표현하자니 Triton 커널 내부는 접근이 불가능한 것이다.

2. 추론용 재귀 스캔(Recursive Scan)의 부재

Mamba의 학습과 추론은 연산 방식이 다르다. 학습 시에는 전체 시퀀스를 한 번에 볼 수 있으므로 병렬 스캔(Parallel Scan) 알고리즘을 사용하여 GPU 병렬성을 극대화한다. 그러나 추론(Generation) 시에는 토큰이 하나씩 들어오므로, 이전 상태(h_{t-1})를 받아 현재 상태(h_t)를 갱신하는 재귀적(Step-wise) 방식이 필요하다.

  • 문제점: 대부분의 오픈소스 Mamba 구현체는 학습용 병렬 스캔에 최적화되어 있다. 이를 그대로 ONNX로 변환하면, 입력 전체가 들어와야만 계산이 가능한 형태가 되어 실시간 스트리밍 추론이 불가능해진다.
  • 해결책: ONNX 변환을 위해서는 mamba_simple.py 등에서 제공하는 추론용 로직(step 함수 등)을 명시적으로 사용하여 그래프를 생성해야 한다. 사용자는 모델의 forward 메서드 대신, 한 단계의 상태 갱신을 수행하는 래퍼(Wrapper) 모듈을 작성하여 이를 export 해야 한다.14

3. 해결 방안: 사용자 정의 연산자(Custom Operator) 구현

Mamba 모델을 성능 저하 없이 ONNX로 변환하는 가장 현실적이고 강력한 방법은 ONNX Runtime(ORT)을 위한 Custom Operator를 구현하는 것이다.

가. 구현 프로세스

  1. PyTorch 심볼릭 함수 정의: torch.autograd.Function을 상속받는 클래스를 만들고, symbolic 정적 메서드를 정의한다. 이 메서드는 ONNX 그래프 상에 domain.SelectiveScan과 같은 이름의 노드를 생성하도록 지시한다.
  2. C++ 커널 작성: ONNX Runtime은 C++로 작성된 백엔드 엔진이다. 따라서 SelectiveScan 노드를 실제로 실행할 C++ 코드가 필요하다. Ort::CustomOp 인터페이스를 상속받아 Compute 메서드 내에 실제 연산 로직(또는 CUDA 커널 호출 코드)을 구현한다.16
  3. 라이브러리 등록: 작성된 C++ 코드를 컴파일하여 공유 라이브러리(.so 또는.dll)로 만들고, Python에서 ONNX Runtime 세션을 초기화할 때 이 라이브러리를 로드하여 사용자 정의 연산자를 등록한다.

나. 순수 PyTorch 구현 (Naive Implementation) 대안

C++ 커널을 직접 작성하는 것이 부담스럽다면, 성능 손실을 감수하고 순수 PyTorch로 selective_scan을 구현하여 변환하는 방법도 있다. mamba.py와 같은 프로젝트에서는 표준 PyTorch 연산(cumsum, exp, mul 등)만을 사용하여 스캔 로직을 구현했다.17

  • 장점: 별도의 C++ 컴파일이나 커스텀 오퍼레이터 등록 없이 표준 ONNX 런타임 어디서나 실행 가능하다.
  • 단점: 메모리 사용량이 증가하고 실행 속도가 현저히 느려진다. 프로덕션 환경보다는 로직 검증이나 저사양 엣지 디바이스 배포에 적합하다.

4. Mamba-2 변환 시의 가중치 형상 불일치 이슈

Mamba-2 아키텍처는 Mamba-1과 달리 텐서 차원 처리 방식이 변경되었다. 특히 RMSNorm 레이어나 헤드 부분에서 체크포인트 파일의 가중치 형상(Shape)과 모델 코드 상의 기대 형상이 일치하지 않는 문제가 자주 보고된다.19

  • 사례: 체크포인트에는 크기의 파라미터가 있는데, Mamba-2 모델 코드는 내부적으로 이를(예: headdim * n_heads)으로 처리하려 할 때 에러가 발생한다.
  • 해결: 변환 스크립트 작성 시, 모델을 로드한 직후 state_dict를 순회하며 불일치가 발생하는 텐서를 찾아 수동으로 reshape 하거나 repeat 하여 형상을 맞춰주는 전처리 코드를 삽입해야 한다.

10.2.5 TensorRT-LLM: 하드웨어 가속의 극한

NVIDIA의 TensorRT-LLM은 vLLM과 달리 NVIDIA GPU 하드웨어에 극도로 최적화된 성능을 제공하는 것을 목표로 한다. Mamba 및 Jamba 지원에 있어서도 TensorRT-LLM은 하드웨어 특성을 활용한 독자적인 최적화 경로를 걷고 있다.

1. Mamba 지원 타임라인과 기능적 특징

TensorRT-LLM은 비교적 이른 시기부터 Mamba-1을 지원하기 시작했으며(v0.8.0 부근), 지속적으로 기능을 확장해 왔다.

  • Mamba-2 및 SSD 지원: 2024년 8월 업데이트(v0.12.0 이후)를 통해 Mamba-2의 핵심인 SSD(Structured State Space Duality)에 대한 지원을 공식화했다.20 이는 Mamba-2의 상태 전파 연산을 텐서 코어(Tensor Core)가 잘 처리할 수 있는 행렬 연산 형태로 변환하여 수행함으로써 처리량을 극대화한다.
  • 텐서 병렬화 (Tensor Parallelism): 초기에는 단일 GPU 지원에 그쳤으나, 현재는 Mamba-2 모델에 대한 텐서 병렬화를 지원한다. 이는 Codestral Mamba 7B와 같은 모델을 여러 GPU에 분산하여 서빙할 때 필수적인 기능이며, 대규모 트래픽 처리를 가능하게 한다.

2. 하이브리드 아키텍처 최적화 전략 (Jamba)

TensorRT-LLM은 Jamba와 같은 하이브리드 모델을 위해 두 가지 핵심 기술을 적용한다.

  • 청크 기반 컨텍스트 (Chunked Context): Mamba는 긴 문맥을 처리할 때 강점을 가지지만, 무작정 긴 시퀀스를 한 번에 처리하면 중간 활성화(Activation) 메모리가 부족해질 수 있다. TensorRT-LLM은 입력 시퀀스를 적절한 크기의 청크(Chunk)로 나누고, 각 청크의 끝에서 계산된 상태(State)를 다음 청크로 넘겨주는 방식(State Passing)을 통해 메모리 사용량을 제어하면서도 긴 문맥을 처리한다.22
  • 커널 퓨전 (Kernel Fusion): TensorRT의 강력한 컴파일러 기능을 활용하여, Mamba의 스캔 연산과 트랜스포머의 MLP, LayerNorm 등을 하나의 커널로 융합한다. 이는 메모리 I/O를 줄이고 GPU 가동률을 높이는 핵심 요인이다.

3. vLLM 대 TensorRT-LLM 성능 및 특징 비교

Mamba 모델 서빙 시 두 엔진의 차이는 명확하다.

특징vLLM (Mamba 지원)TensorRT-LLM (Mamba 지원)
핵심 철학유연성, 쉬운 배포, PagedAttention하드웨어 극한 성능, 컴파일 최적화
하이브리드 지원Jamba 지원 (메모리 패딩 이슈 존재)Jamba 지원 (청크 처리 최적화)
멀티 GPU지원 (구현체에 따라 효율 차이)강력한 TP 지원 (대규모 모델 유리)
메모리 관리통합 할당자(v1)로 개선 중정적 할당 및 엔진 빌드 시 결정
프리픽스 캐싱미지원 (구조적 어려움)아키텍처상 제한적
배포 난이도낮음 (Python 기반, HF 모델 로드)높음 (엔진 빌드 과정 필요)
ONNX 필요성없음 (자체 실행)없음 (자체 포맷 변환)

벤치마크 결과들을 종합하면, vLLM은 사용 편의성과 PagedAttention을 통한 동시성 처리에 강점이 있으나, Mamba 모델에 대해서는 아직 트랜스포머만큼의 최적화(프리픽스 캐싱 등)가 이루어지지 않았다.23 반면, TensorRT-LLM은 엔진 빌드라는 번거로운 과정이 있지만, 일단 빌드된 엔진은 H100 등 최신 하드웨어에서 vLLM 대비 높은 처리량(Throughput)을 보여준다.24

10.2.6 사례 연구: Jamba 모델의 배포 파이프라인 구축

실제 프로덕션 환경에서 AI21 Labs의 Jamba 1.5 모델을 배포한다고 가정할 때, 엔지니어는 다음과 같은 파이프라인을 고려해야 한다.

  1. 모델 준비: Hugging Face에서 ai21labs/AI21-Jamba-1.5-Large 모델을 다운로드한다. 이 모델은 전문가(Experts) 레이어를 포함하고 있어 메모리 요구량이 크다.
  2. 엔진 선택:
    • 빠른 프로토타이핑/연구: vLLM을 선택한다. pip install vllmvllm serve ai21labs/AI21-Jamba-1.5-Large --tensor-parallel-size 8 명령만으로 즉시 API 서버를 띄울 수 있다. 이때 --quantization experts_int8 옵션을 주어 전문가 레이어를 양자화하면 메모리를 절약할 수 있다.11
    • 대규모 상용 서비스: TensorRT-LLM을 선택한다. 먼저 TensorRT-LLM 도커 컨테이너를 실행하고, 체크포인트 변환 스크립트를 통해 Jamba 모델을 TensorRT 엔진 포맷으로 빌드한다. 이 과정에서 max_batch_size, max_input_len 등을 서비스 요구사항에 맞춰 정적으로 설정해야 한다. 빌드된 엔진은 Triton Inference Server에 얹어 배포한다.3
  3. 트러블슈팅:
    • vLLM 사용 시 OOM(Out of Memory)가 발생하면, block_size 설정을 확인하거나 max_num_seqs를 줄인다. Jamba의 경우 Mamba 상태 저장을 위해 예상보다 많은 메모리를 패딩으로 소모할 수 있음을 기억해야 한다.9
    • TensorRT-LLM 빌드 실패 시, 모델의 특정 레이어(특히 MoE 라우터나 Mamba 믹서)가 지원되는 버전인지 확인하고, 필요하다면 NVIDIA 개발자 포럼이나 GitHub 이슈를 통해 패치를 확인한다.

10.2.7 결론 및 향후 전망

Mamba와 Jamba로 대표되는 비(非) 트랜스포머 아키텍처는 서빙 인프라에 새로운 패러다임을 요구하고 있다. vLLM은 v1 업데이트를 통해 이질적인 메모리 요구사항을 가진 하이브리드 모델을 통합하려는 아키텍처적 유연성을 보여주고 있으며, TensorRT-LLM은 하드웨어 제조사의 이점을 살려 SSM의 연산 효율을 극한으로 끌어올리고 있다.

향후 서빙 기술은 **‘상태(State) 관리의 고도화’**로 수렴될 것이다. Mamba의 상태를 효율적으로 캐싱하고 공유할 수 있는 알고리즘이 개발된다면(마치 트랜스포머의 프리픽스 캐싱처럼), SSM 기반 모델은 긴 문맥 처리뿐만 아니라 다중 턴 대화나 에이전트 작업에서도 트랜스포머를 능가하는 효율을 보여줄 것이다. 또한, ONNX와 같은 표준 포맷 역시 커뮤니티의 노력으로 Selective Scan과 같은 커스텀 연산자를 표준화하거나 더 쉽게 통합하는 방향으로 발전할 것으로 예상된다. 엔지니어들은 이러한 변화의 흐름 속에서 각 엔진의 장단점을 명확히 이해하고, 서비스의 특성(지연 시간 vs 처리량, 유연성 vs 성능)에 맞는 최적의 도구를 선택해야 할 것이다.


표 10.2.1 서빙 엔진별 Mamba/Jamba 지원 기능 상세 비교

기능 (Feature)vLLM (v0.6.3+ / v1 Preview)TensorRT-LLM (v0.12+)비고
Mamba-1 지원지원함 (MambaForCausalLM)지원함
Mamba-2 지원지원함 (일부 기능 제한적)지원함 (SSD 가속 포함)TRT-LLM은 TP 지원 및 SSD 최적화 우세
Jamba (Hybrid) 지원지원함 (MoE 양자화 포함)지원함vLLM은 메모리 패딩 이슈 존재
텐서 병렬화 (TP)지원함강력 지원대규모 클러스터 서빙 시 TRT-LLM 권장
프리픽스 캐싱미지원 (기술적 난제)아키텍처상 제한적챗봇 등 공유 컨텍스트 환경에서 약점
메모리 할당 방식PagedAttention (동적, 블록 기반)정적 할당 및 Paged KV (제한적)vLLM v1에서 통합 할당자로 개선 중
ONNX 변환 용이성해당 없음 (자체 실행 엔진)해당 없음 (자체 엔진 포맷 사용)ONNX 변환은 별도 툴체인(ORT) 영역

표 10.2.2 Mamba ONNX 변환 시 주요 에러 및 기술적 해결책

에러 메시지 / 증상원인 (Root Cause)해결 방안 (Solution)
TracerWarning: Converting a tensor to a Python integer...스캔 루프 내 동적 제어 흐름을 Tracer가 상수로 고정함Scripting 방식 사용 또는 루프 구조 재작성, torch.onnx.exportdynamic_axes 설정 점검
Unknown operator: SelectiveScanselective_scan이 표준 ONNX Opset에 없음Custom Operator(C++) 구현 및 ORT 등록, 또는 순수 PyTorch 구현으로 대체(성능 저하 감수)
Weight shape mismatch (Mamba-2)Mamba-2 RMSNorm 등의 가중치 형상이 코드와 체크포인트 간 불일치변환 스크립트 전처리 단계에서 state_dict의 가중치 텐서를 reshape 또는 repeat 처리
Triton kernel not supportedtorch.onnx.export가 컴파일된 Triton/CUDA 커널 내부를 추적 불가Python 레벨의 ‘Recursive Scan’ (추론용) 구현체로 교체하여 Export 수행

참고 자료

  1. Zamba: A Compact 7B SSM Hybrid Model - arXiv, 12월 26, 2025에 액세스, https://arxiv.org/html/2405.16712v1
  2. Mamba vs Transformers: Efficiency, Scale, and the Future of AI - Michiel Horstman - Medium, 12월 26, 2025에 액세스, https://michielh.medium.com/mamba-vs-transformers-efficiency-scale-and-the-future-of-ai-d7a8dedb4018
  3. Jamba 1.5 LLMs Leverage Hybrid Architecture to Deliver Superior Reasoning and Long Context Handling - NVIDIA Developer, 12월 26, 2025에 액세스, https://developer.nvidia.com/blog/jamba-1-5-llms-leverage-hybrid-architecture-to-deliver-superior-reasoning-and-long-context-handling/
  4. Paged Attention - vLLM, 12월 26, 2025에 액세스, https://docs.vllm.ai/en/stable/design/paged_attention/
  5. [RFC]: Native support for Mamba, SSM, and hybrid transformer models in vLLM V1 #17140, 12월 26, 2025에 액세스, https://github.com/vllm-project/vllm/issues/17140
  6. [RFC]: Hybrid Memory Allocator · Issue #11382 · vllm-project/vllm - GitHub, 12월 26, 2025에 액세스, https://github.com/vllm-project/vllm/issues/11382
  7. single_type_kv_cache_manager - vLLM, 12월 26, 2025에 액세스, https://docs.vllm.ai/en/v0.10.2/api/vllm/v1/core/single_type_kv_cache_manager.html
  8. Hybrid Models as First-Class Citizens in vLLM - PyTorch, 12월 26, 2025에 액세스, https://pytorch.org/blog/hybrid-models-as-first-class-citizens-in-vllm/
  9. Hybrid KV Cache Manager - vLLM, 12월 26, 2025에 액세스, https://docs.vllm.ai/en/latest/design/hybrid_kv_cache_manager/
  10. Supported Models - vLLM, 12월 26, 2025에 액세스, https://docs.vllm.ai/en/latest/models/supported_models.html
  11. vLLM - AI21 Docs, 12월 26, 2025에 액세스, https://docs.ai21.com/docs/vllm
  12. mamba2 - vLLM, 12월 26, 2025에 액세스, https://docs.vllm.ai/en/latest/api/vllm/model_executor/models/mamba2/
  13. [Bug]: Speculative decoding support for mamba models · Issue #30114 - GitHub, 12월 26, 2025에 액세스, https://github.com/vllm-project/vllm/issues/30114
  14. Generating onnx file for the inference of Mamba? · Issue #200 - GitHub, 12월 26, 2025에 액세스, https://github.com/state-spaces/mamba/issues/200
  15. There was a problem converting the model to onnx format - PyTorch Forums, 12월 26, 2025에 액세스, https://discuss.pytorch.org/t/there-was-a-problem-converting-the-model-to-onnx-format/151356
  16. Custom operators | onnxruntime, 12월 26, 2025에 액세스, https://onnxruntime.ai/docs/reference/operators/add-custom-op.html
  17. mamba.py/README.md at main · alxndrTL/mamba.py · GitHub, 12월 26, 2025에 액세스, https://github.com/alxndrTL/mamba.py/blob/main/README.md
  18. Mamba - Hugging Face, 12월 26, 2025에 액세스, https://huggingface.co/docs/transformers/en/model_doc/mamba
  19. ONNX Export Failure for state-spaces/mamba2-370m: Mamba2RMSNorm Weight Shape Mismatch · Issue #751 - GitHub, 12월 26, 2025에 액세스, https://github.com/state-spaces/mamba/issues/751
  20. Update on the development branch · NVIDIA TensorRT-LLM · Discussion #2131 - GitHub, 12월 26, 2025에 액세스, https://github.com/NVIDIA/TensorRT-LLM/discussions/2131
  21. TensorRT-LLM 0.14.0 Release #2403 - GitHub, 12월 26, 2025에 액세스, https://github.com/NVIDIA/TensorRT-LLM/discussions/2403
  22. Revolutionizing Code Completion with Codestral Mamba, the Next-Gen Coding LLM, 12월 26, 2025에 액세스, https://developer.nvidia.com/blog/revolutionizing-code-completion-with-codestral-mamba-the-next-gen-coding-llm/
  23. vLLM vs TensorRT-LLM: Key differences, performance, and how to run them - Northflank, 12월 26, 2025에 액세스, https://northflank.com/blog/vllm-vs-tensorrt-llm-and-how-to-run-them
  24. Why is vLLM Outperforming TensorRT-LLM (Nvidia’s deployment library)? My Shocking Benchmarks on GPT-OSS-120B with H100 : r/LocalLLaMA - Reddit, 12월 26, 2025에 액세스, https://www.reddit.com/r/LocalLLaMA/comments/1oyawkl/why_is_vllm_outperforming_tensorrtllm_nvidias/
  25. Deploy NeMo Models by Exporting TensorRT-LLM - NVIDIA Documentation, 12월 26, 2025에 액세스, https://docs.nvidia.com/nemo-framework/user-guide/24.12/deployment/llm/optimized/tensorrt_llm.html